#!/usr/bin/env python
# author:AnFany
# datetime:2024/07/11 15:52
from ultralytics import YOLO
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image
import pandas as pd
import os
import argparse
import gc

plt.rcParams['font.sans-serif'] = ['SimHei']

# 定义一个设备黑名单，在名单中的不跑预警
device_black_lst = {'waterflow': [],
                    'sewage': [],
                    'float':[]
                    }
# 判断一个点是不是在一个不规则的多边形状内
def point_in_poly(xy, poly):
    x, y = xy
    num = len(poly)
    i, j = 0, num - 1
    c = False
    for i in range(num):
        if (((poly[i][1] > y) != (poly[j][1] > y)) and
                (x < poly[i][0] + (poly[j][0] - poly[i][0]) * (y - poly[i][1]) / (poly[j][1] - poly[i][1]))):
            c = not c
        j = i
    return c


# 预测图片
def tupian_yujing(figpath, csvpath, warnsave, modelpath, device_name,
                  classname=None, colorname=None, warnnumber=None, beizhuname=None):
    """
    :param figpath: 待识别的图片绝对路径
    :param csvpath: 待识别图片站点的框的绝对路径
    :param warnsave: 带有识别检测框的存储的绝对路径
    :param modelpath: 训练好的模型
    :parm devicename:设备名称
    """
    # 报警类型
    WarnLei = []
    sign = 1
    # 加载模型
    if classname is None:
        classname = {0: 'waterflow', 1: 'sewage', 2: 'float'}
    # 类别框以及类别文字的颜色
    if colorname is None:
        colorname = {'waterflow': '#23ee56', 'sewage': 'blue', 'float':'red'}
    # 各个类别的置信度阈值
    if warnnumber is None:
        warnnumber = {'waterflow': 0.5, 'sewage': 0.5, 'float': 0.5}

    # 报警的备注
    if beizhuname is None:
        beizhuname = {'waterflow': '排水', 'sewage': '污水', 'float': '漂浮物'}

    classidlist = {'waterflow': 0, 'sewage': 1, 'float': 2}
    box_lst = {}

    # 读取csv框中的多边形
    # 多边形数据
    
    if csvpath is not None:
        try:
            data_file = pd.read_csv(csvpath)
        except:
            data_file = pd.read_csv(csvpath, encoding='gbk')
        # 使用iterrows()遍历行
        for index, row in data_file.iterrows():
            typename = row['type']
            if typename not in box_lst:
                box_lst[typename] = []
            data = eval(row['region_shape_attributes'])
            x_p = data['all_points_x']
            y_p = data['all_points_y']
            xypp = list(zip(x_p, y_p))
            box_lst[typename].append(xypp)


    # 加载模型
    #model = YOLO(modelpath)
    # 预测图片
    result = modelpath.predict(figpath, save=False, device='cpu')
    # 获得类别
    leibie = [int(k) for k in result[0].boxes.cls.tolist()]
    # 获得置信度
    zhixindu = [float('%.4f' % k) for k in result[0].boxes.conf.tolist()]
    # 获得框
    xyxy = result[0].boxes.xyxy.tolist()

    
    
    # 未经过csv框限制和conf置信度限制的预警图片
    # 打开图片
    image = Image.open(figpath)
    # 创建一个新的绘图
    fig1, ax1 = plt.subplots()
    # 去除坐标轴
    ax1.axis('off')
    # 显示图像
    ax1.imshow(image)
    # 原始报警画框
    for kin, kcl in enumerate(leibie):
        # 类别
        leibie_name = classname[int(kcl)]
        # 获得置信度
        coef_float = zhixindu[kin]

        zuox, zuoy, youx, youy = xyxy[kin]

        rectangle = patches.Rectangle((zuox, zuoy), youx - zuox, youy - zuoy,
                                      linewidth=1.6, edgecolor=colorname[leibie_name],
                                      facecolor='none')
        # 添加矩形框到图像
        ax1.add_patch(rectangle)
        # 找准位置添加文字，在框的左上角的上面
        # 写的内容
        baojingliexing = beizhuname[leibie_name]
        str_name = '{}:{}%'.format(baojingliexing, '%.2f' % (coef_float * 100))
        ax1.text(zuox - 5, zuoy - 22, str_name, color='w', fontsize=6,
                fontweight='bold', bbox=dict(facecolor=colorname[leibie_name], alpha=0.5, pad=0.1))
        
        # 获得图片的名称
        figname = 'no_csv_' + os.path.basename(figpath)
        # 保存图片
        plt.savefig(os.path.join(warnsave, figname), dpi=300, bbox_inches='tight')


    # 经过csv框限制和conf置信度限制的预警图片
    # 自己画框
    fig, ax = plt.subplots()
    # 去除坐标轴
    ax.axis('off')
    # 显示图像
    ax.imshow(image)
    for kin, kcl in enumerate(leibie):
        leibie_name = classname[int(kcl)]
        # 获得置信度
        coef_float = zhixindu[kin]
        # 大于就报警
        if coef_float > warnnumber[leibie_name]:
            # 开始画框
            zuox, zuoy, youx, youy = xyxy[kin]
            # 增加是否在站点筛选框中的逻辑
            # 获得四个点
            sign = 0
            in_count = 0
            for kj in [[zuox, zuoy], [zuox, youy], [youx, zuoy], [youx, youy]]:
                # 有一个在筛选框内，就不绘制
                if csvpath is not None and leibie_name in box_lst:

                    for box_xypp in box_lst[leibie_name]:
                        if point_in_poly(kj, box_xypp):
                            in_count += 1
                            break
                    if in_count == 4:
                        sign = 1

            # 4个坐标点都在报警框内才报警
            if sign and device_name not in device_black_lst[leibie_name]:
                rectangle = patches.Rectangle((zuox, zuoy), youx - zuox, youy - zuoy,
                                              linewidth=1.6, edgecolor=colorname[leibie_name],
                                              facecolor='none')
                # 添加矩形框到图像
                ax.add_patch(rectangle)
                # 找准位置添加文字，在框的左上角的上面
                # 写的内容
                baojingliexing = beizhuname[leibie_name]
                clsid = classidlist[leibie_name]
                str_name = '{}:{}%'.format(baojingliexing, '%.2f' % (coef_float * 100))
                ax.text(zuox - 5, zuoy - 22, str_name, color='w', fontsize=6,
                        fontweight='bold', bbox=dict(facecolor=colorname[leibie_name], alpha=0.5, pad=0.1))
                # 储存报警类型
                if clsid not in WarnLei:
                    WarnLei.append(clsid)

    # 显示图像和矩形框
    if len(WarnLei):
        # 获得图片的名称
        figname = os.path.basename(figpath)
        # 保存图片
        plt.savefig(os.path.join(warnsave, figname), dpi=300, bbox_inches='tight')
    plt.close('all')
    # 强制垃圾回收
    gc.collect()
    return WarnLei


if __name__ == '__main__':
    # 创建ArgumentParser对象：
    parser = argparse.ArgumentParser()

    # 添加命令行参数：
    parser.add_argument('--figpath')
    parser.add_argument('--csvpath')
    parser.add_argument('--warnsave')
    parser.add_argument('--modelpath')
    parser.add_argument('--device_name')

    # 解析命令行参数：
    args = parser.parse_args()

    example = tupian_yujing(figpath=args.figpath,
                            csvpath=args.csvpath,
                            warnsave=args.warnsave,
                            modelpath=args.modelpath,
                            device_name=args.device_name)
    print(example)
